// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package qunar.tc.decompiler.modules.code;

import qunar.tc.decompiler.code.CodeConstants;
import qunar.tc.decompiler.code.Instruction;
import qunar.tc.decompiler.code.InstructionSequence;
import qunar.tc.decompiler.code.SimpleInstructionSequence;
import qunar.tc.decompiler.code.cfg.BasicBlock;
import qunar.tc.decompiler.code.cfg.ControlFlowGraph;
import qunar.tc.decompiler.code.cfg.ExceptionRangeCFG;
import qunar.tc.decompiler.main.DecompilerContext;
import qunar.tc.decompiler.main.extern.IFernflowerPreferences;

import java.util.*;

public class DeadCodeHelper {

    public static void removeDeadBlocks(ControlFlowGraph graph) {

        LinkedList<BasicBlock> stack = new LinkedList<>();
        HashSet<BasicBlock> setStacked = new HashSet<>();

        stack.add(graph.getFirst());
        setStacked.add(graph.getFirst());

        while (!stack.isEmpty()) {
            BasicBlock block = stack.removeFirst();

            List<BasicBlock> lstSuccs = new ArrayList<>(block.getSuccs());
            lstSuccs.addAll(block.getSuccExceptions());

            for (BasicBlock succ : lstSuccs) {
                if (!setStacked.contains(succ)) {
                    stack.add(succ);
                    setStacked.add(succ);
                }
            }
        }

        HashSet<BasicBlock> setAllBlocks = new HashSet<>(graph.getBlocks());
        setAllBlocks.removeAll(setStacked);

        for (BasicBlock block : setAllBlocks) {
            graph.removeBlock(block);
        }
    }

    public static void removeEmptyBlocks(ControlFlowGraph graph) {

        List<BasicBlock> blocks = graph.getBlocks();

        boolean cont;
        do {
            cont = false;

            for (int i = blocks.size() - 1; i >= 0; i--) {
                BasicBlock block = blocks.get(i);

                if (removeEmptyBlock(graph, block, false)) {
                    cont = true;
                    break;
                }
            }
        }
        while (cont);
    }

    private static boolean removeEmptyBlock(ControlFlowGraph graph, BasicBlock block, boolean merging) {

        boolean deletedRanges = false;

        if (block.getSeq().isEmpty()) {

            if (block.getSuccs().size() > 1) {
                if (block.getPreds().size() > 1) {
                    // ambiguous block
                    throw new RuntimeException("ERROR: empty block with multiple predecessors and successors found");
                } else if (!merging) {
                    throw new RuntimeException("ERROR: empty block with multiple successors found");
                }
            }

            HashSet<BasicBlock> setExits = new HashSet<>(graph.getLast().getPreds());

            if (block.getPredExceptions().isEmpty() &&
                    (!setExits.contains(block) || block.getPreds().size() == 1)) {

                if (setExits.contains(block)) {
                    BasicBlock pred = block.getPreds().get(0);

                    // FIXME: flag in the basic block
                    if (pred.getSuccs().size() != 1 || (!pred.getSeq().isEmpty()
                            && pred.getSeq().getLastInstr().group == CodeConstants.GROUP_SWITCH)) {
                        return false;
                    }
                }

                HashSet<BasicBlock> setPreds = new HashSet<>(block.getPreds());
                HashSet<BasicBlock> setSuccs = new HashSet<>(block.getSuccs());

                // collect common exception ranges of predecessors and successors
                HashSet<BasicBlock> setCommonExceptionHandlers = null;
                for (int i = 0; i < 2; ++i) {
                    for (BasicBlock pred : i == 0 ? setPreds : setSuccs) {
                        if (setCommonExceptionHandlers == null) {
                            setCommonExceptionHandlers = new HashSet<>(pred.getSuccExceptions());
                        } else {
                            setCommonExceptionHandlers.retainAll(pred.getSuccExceptions());
                        }
                    }
                }

                // check the block to be in each of the common ranges
                if (setCommonExceptionHandlers != null && !setCommonExceptionHandlers.isEmpty()) {
                    for (BasicBlock handler : setCommonExceptionHandlers) {
                        if (!block.getSuccExceptions().contains(handler)) {
                            return false;
                        }
                    }
                }

                // remove ranges consisting of this one block
                List<ExceptionRangeCFG> lstRanges = graph.getExceptions();
                for (int i = lstRanges.size() - 1; i >= 0; i--) {
                    ExceptionRangeCFG range = lstRanges.get(i);
                    List<BasicBlock> lst = range.getProtectedRange();

                    if (lst.size() == 1 && lst.get(0) == block) {
                        if (DecompilerContext.getOption(IFernflowerPreferences.REMOVE_EMPTY_RANGES)) {
                            block.removeSuccessorException(range.getHandler());
                            lstRanges.remove(i);

                            deletedRanges = true;
                        } else {
                            return false;
                        }
                    }
                }


                // connect remaining nodes
                if (merging) {
                    BasicBlock pred = block.getPreds().get(0);
                    pred.removeSuccessor(block);

                    List<BasicBlock> lstSuccs = new ArrayList<>(block.getSuccs());
                    for (BasicBlock succ : lstSuccs) {
                        block.removeSuccessor(succ);
                        pred.addSuccessor(succ);
                    }
                } else {
                    for (BasicBlock pred : setPreds) {
                        for (BasicBlock succ : setSuccs) {
                            pred.replaceSuccessor(block, succ);
                        }
                    }
                }

                // finally exit edges
                Set<BasicBlock> setFinallyExits = graph.getFinallyExits();
                if (setFinallyExits.contains(block)) {
                    setFinallyExits.remove(block);
                    setFinallyExits.add(setPreds.iterator().next());
                }

                // replace first if necessary
                if (graph.getFirst() == block) {
                    if (setSuccs.size() != 1) {
                        throw new RuntimeException("multiple or no entry blocks!");
                    } else {
                        graph.setFirst(setSuccs.iterator().next());
                    }
                }

                // remove this block
                graph.removeBlock(block);

                if (deletedRanges) {
                    removeDeadBlocks(graph);
                }
            }
        }

        return deletedRanges;
    }


    public static boolean isDominator(ControlFlowGraph graph, BasicBlock block, BasicBlock dom) {

        HashSet<BasicBlock> marked = new HashSet<>();

        if (block == dom) {
            return true;
        }

        LinkedList<BasicBlock> lstNodes = new LinkedList<>();
        lstNodes.add(block);

        while (!lstNodes.isEmpty()) {

            BasicBlock node = lstNodes.remove(0);
            if (marked.contains(node)) {
                continue;
            } else {
                marked.add(node);
            }

            if (node == graph.getFirst()) {
                return false;
            }

            for (int i = 0; i < node.getPreds().size(); i++) {
                BasicBlock pred = node.getPreds().get(i);
                if (pred != dom && !marked.contains(pred)) {
                    lstNodes.add(pred);
                }
            }

            for (int i = 0; i < node.getPredExceptions().size(); i++) {
                BasicBlock pred = node.getPredExceptions().get(i);
                if (pred != dom && !marked.contains(pred)) {
                    lstNodes.add(pred);
                }
            }
        }

        return true;
    }

    public static void removeGotos(ControlFlowGraph graph) {

        for (BasicBlock block : graph.getBlocks()) {
            Instruction instr = block.getLastInstruction();

            if (instr != null && instr.opcode == CodeConstants.opc_goto) {
                block.getSeq().removeLast();
            }
        }

        removeEmptyBlocks(graph);
    }

    public static void connectDummyExitBlock(ControlFlowGraph graph) {

        BasicBlock exit = graph.getLast();
        for (BasicBlock block : new HashSet<>(exit.getPreds())) {
            exit.removePredecessor(block);
            block.addSuccessor(exit);
        }
    }

    public static void extendSynchronizedRangeToMonitorexit(ControlFlowGraph graph) {

        while (true) {

            boolean range_extended = false;

            for (ExceptionRangeCFG range : graph.getExceptions()) {

                Set<BasicBlock> setPreds = new HashSet<>();
                for (BasicBlock block : range.getProtectedRange()) {
                    setPreds.addAll(block.getPreds());
                }
                setPreds.removeAll(range.getProtectedRange());

                if (setPreds.size() != 1) {
                    continue; // multiple predecessors, obfuscated range
                }

                BasicBlock predBlock = setPreds.iterator().next();
                InstructionSequence predSeq = predBlock.getSeq();
                if (predSeq.isEmpty() || predSeq.getLastInstr().opcode != CodeConstants.opc_monitorenter) {
                    continue; // not a synchronized range
                }

                boolean monitorexit_in_range = false;
                Set<BasicBlock> setProtectedBlocks = new HashSet<>();
                setProtectedBlocks.addAll(range.getProtectedRange());
                setProtectedBlocks.add(range.getHandler());

                for (BasicBlock block : setProtectedBlocks) {
                    InstructionSequence blockSeq = block.getSeq();
                    for (int i = 0; i < blockSeq.length(); i++) {
                        if (blockSeq.getInstr(i).opcode == CodeConstants.opc_monitorexit) {
                            monitorexit_in_range = true;
                            break;
                        }
                    }
                    if (monitorexit_in_range) {
                        break;
                    }
                }

                if (monitorexit_in_range) {
                    continue; // protected range already contains monitorexit
                }

                Set<BasicBlock> setSuccs = new HashSet<>();
                for (BasicBlock block : range.getProtectedRange()) {
                    setSuccs.addAll(block.getSuccs());
                }
                setSuccs.removeAll(range.getProtectedRange());

                if (setSuccs.size() != 1) {
                    continue; // non-unique successor
                }

                BasicBlock succBlock = setSuccs.iterator().next();
                InstructionSequence succSeq = succBlock.getSeq();

                int succ_monitorexit_index = -1;
                for (int i = 0; i < succSeq.length(); i++) {
                    if (succSeq.getInstr(i).opcode == CodeConstants.opc_monitorexit) {
                        succ_monitorexit_index = i;
                        break;
                    }
                }

                if (succ_monitorexit_index < 0) {
                    continue; // monitorexit not found in the single successor block
                }

                BasicBlock handlerBlock = range.getHandler();
                if (handlerBlock.getSuccs().size() != 1) {
                    continue; // non-unique handler successor
                }
                BasicBlock succHandler = handlerBlock.getSuccs().get(0);
                InstructionSequence succHandlerSeq = succHandler.getSeq();
                if (succHandlerSeq.isEmpty() || succHandlerSeq.getLastInstr().opcode != CodeConstants.opc_athrow) {
                    continue; // not a standard synchronized range
                }

                int handler_monitorexit_index = -1;
                for (int i = 0; i < succHandlerSeq.length(); i++) {
                    if (succHandlerSeq.getInstr(i).opcode == CodeConstants.opc_monitorexit) {
                        handler_monitorexit_index = i;
                        break;
                    }
                }

                if (handler_monitorexit_index < 0) {
                    continue; // monitorexit not found in the handler successor block
                }

                // checks successful, prerequisites satisfied, now extend the range
                if (succ_monitorexit_index < succSeq.length() - 1) { // split block

                    SimpleInstructionSequence seq = new SimpleInstructionSequence();
                    for (int counter = 0; counter < succ_monitorexit_index; counter++) {
                        seq.addInstruction(succSeq.getInstr(0), -1);
                        succSeq.removeInstruction(0);
                    }

                    // build a separate block
                    BasicBlock newblock = new BasicBlock(++graph.last_id);
                    newblock.setSeq(seq);

                    // insert new block
                    for (BasicBlock block : succBlock.getPreds()) {
                        block.replaceSuccessor(succBlock, newblock);
                    }

                    newblock.addSuccessor(succBlock);
                    graph.getBlocks().addWithKey(newblock, newblock.id);

                    succBlock = newblock;
                }

                // copy exception edges and extend protected ranges (successor block)
                BasicBlock rangeExitBlock = succBlock.getPreds().get(0);
                for (int j = 0; j < rangeExitBlock.getSuccExceptions().size(); j++) {
                    BasicBlock hd = rangeExitBlock.getSuccExceptions().get(j);
                    succBlock.addSuccessorException(hd);

                    ExceptionRangeCFG rng = graph.getExceptionRange(hd, rangeExitBlock);
                    rng.getProtectedRange().add(succBlock);
                }

                // copy instructions (handler successor block)
                InstructionSequence handlerSeq = handlerBlock.getSeq();
                for (int counter = 0; counter < handler_monitorexit_index; counter++) {
                    handlerSeq.addInstruction(succHandlerSeq.getInstr(0), -1);
                    succHandlerSeq.removeInstruction(0);
                }

                range_extended = true;
                break;
            }

            if (!range_extended) {
                break;
            }
        }

    }


    public static void incorporateValueReturns(ControlFlowGraph graph) {

        for (BasicBlock block : graph.getBlocks()) {
            InstructionSequence seq = block.getSeq();

            int len = seq.length();
            if (len > 0 && len < 3) {

                boolean ok = false;

                if (seq.getLastInstr().opcode >= CodeConstants.opc_ireturn && seq.getLastInstr().opcode <= CodeConstants.opc_return) {
                    if (len == 1) {
                        ok = true;
                    } else if (seq.getLastInstr().opcode != CodeConstants.opc_return) {
                        switch (seq.getInstr(0).opcode) {
                            case CodeConstants.opc_iload:
                            case CodeConstants.opc_lload:
                            case CodeConstants.opc_fload:
                            case CodeConstants.opc_dload:
                            case CodeConstants.opc_aload:
                            case CodeConstants.opc_aconst_null:
                            case CodeConstants.opc_bipush:
                            case CodeConstants.opc_sipush:
                            case CodeConstants.opc_lconst_0:
                            case CodeConstants.opc_lconst_1:
                            case CodeConstants.opc_fconst_0:
                            case CodeConstants.opc_fconst_1:
                            case CodeConstants.opc_fconst_2:
                            case CodeConstants.opc_dconst_0:
                            case CodeConstants.opc_dconst_1:
                            case CodeConstants.opc_ldc:
                            case CodeConstants.opc_ldc_w:
                            case CodeConstants.opc_ldc2_w:
                                ok = true;
                        }
                    }
                }

                if (ok) {

                    if (!block.getPreds().isEmpty()) {

                        HashSet<BasicBlock> setPredHandlersUnion = new HashSet<>();
                        HashSet<BasicBlock> setPredHandlersIntersection = new HashSet<>();

                        boolean firstpred = true;
                        for (BasicBlock pred : block.getPreds()) {
                            if (firstpred) {
                                setPredHandlersIntersection.addAll(pred.getSuccExceptions());
                                firstpred = false;
                            } else {
                                setPredHandlersIntersection.retainAll(pred.getSuccExceptions());
                            }

                            setPredHandlersUnion.addAll(pred.getSuccExceptions());
                        }

                        // add exception ranges from predecessors
                        setPredHandlersIntersection.removeAll(block.getSuccExceptions());
                        BasicBlock predecessor = block.getPreds().get(0);

                        for (BasicBlock handler : setPredHandlersIntersection) {
                            ExceptionRangeCFG range = graph.getExceptionRange(handler, predecessor);

                            range.getProtectedRange().add(block);
                            block.addSuccessorException(handler);
                        }

                        // remove redundant ranges
                        HashSet<BasicBlock> setRangesToBeRemoved = new HashSet<>(block.getSuccExceptions());
                        setRangesToBeRemoved.removeAll(setPredHandlersUnion);

                        for (BasicBlock handler : setRangesToBeRemoved) {
                            ExceptionRangeCFG range = graph.getExceptionRange(handler, block);

                            if (range.getProtectedRange().size() > 1) {
                                range.getProtectedRange().remove(block);
                                block.removeSuccessorException(handler);
                            }
                        }
                    }


                    if (block.getPreds().size() == 1 && block.getPredExceptions().isEmpty()) {

                        BasicBlock bpred = block.getPreds().get(0);
                        if (bpred.getSuccs().size() == 1) {

                            // add exception ranges of predecessor
                            for (BasicBlock succ : bpred.getSuccExceptions()) {
                                if (!block.getSuccExceptions().contains(succ)) {
                                    ExceptionRangeCFG range = graph.getExceptionRange(succ, bpred);

                                    range.getProtectedRange().add(block);
                                    block.addSuccessorException(succ);
                                }
                            }

                            // remove superfluous ranges from successors
                            for (BasicBlock succ : new HashSet<>(block.getSuccExceptions())) {
                                if (!bpred.getSuccExceptions().contains(succ)) {
                                    ExceptionRangeCFG range = graph.getExceptionRange(succ, block);

                                    if (range.getProtectedRange().size() > 1) {
                                        range.getProtectedRange().remove(block);
                                        block.removeSuccessorException(succ);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }


    public static void mergeBasicBlocks(ControlFlowGraph graph) {

        while (true) {

            boolean merged = false;

            for (BasicBlock block : graph.getBlocks()) {

                InstructionSequence seq = block.getSeq();

                if (block.getSuccs().size() == 1) {
                    BasicBlock next = block.getSuccs().get(0);

                    if (next != graph.getLast() && (seq.isEmpty() || seq.getLastInstr().group != CodeConstants.GROUP_SWITCH)) {

                        if (next.getPreds().size() == 1 && next.getPredExceptions().isEmpty()
                                && next != graph.getFirst()) {
                            // TODO: implement a dummy start block
                            boolean sameRanges = true;
                            for (ExceptionRangeCFG range : graph.getExceptions()) {
                                if (range.getProtectedRange().contains(block) ^
                                        range.getProtectedRange().contains(next)) {
                                    sameRanges = false;
                                    break;
                                }
                            }

                            if (sameRanges) {
                                seq.addSequence(next.getSeq());
                                block.getInstrOldOffsets().addAll(next.getInstrOldOffsets());
                                next.getSeq().clear();

                                removeEmptyBlock(graph, next, true);

                                merged = true;
                                break;
                            }
                        }
                    }
                }
            }

            if (!merged) {
                break;
            }
        }
    }
}
